import os
import logging
import numpy as np
import pandas as pd
import tulipy as ti
import talib as ta
import psycopg2.extras
import backtrader as bt
from datetime import datetime
import plotly.express as px
import plotly.graph_objects as go
from Units import QUnits, CustomError

logging.basicConfig(format="%(asctime)s %(filename)s(line:%(lineno)d) [%(levelname)s] : %(message)s", datefmt="%Y-%M-%d %H:%M:%S", level=logging.DEBUG)

class StockAnalysis:
    def __init__(self):
        self.units = QUnits()
        self.setting = self.units.XML2ObjFromFile('./setting.xml').setting

    def getStockData(self, code_ids, category):
        conn = self.units.dbConnect(self.setting.pgsql)
        stockSecurityData = dict()
        try:
            cursor = conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
            cursor.execute("SELECT * FROM trading.stock_view  WHERE stock_code IN %s" % str(tuple(code_ids)))
            for row in cursor.fetchall():
                stockSecurityData[row['stock_code']] = dict()
                stockSecurityData[row['stock_code']]['info'] = row
                stockSecurityData[row['stock_code']]['data'] = self.units.getStockSecurityBars(category, row['market'], row['stock_code'], 1, 300, True)
                stockSecurityData[row['stock_code']]['data'].index = pd.to_datetime(stockSecurityData[row['stock_code']]['data'].datetime)
                stockSecurityData[row['stock_code']]['data']['volume'] = stockSecurityData[row['stock_code']]['data'].vol
                stockSecurityData[row['stock_code']]['data']['openinterest'] = 0
        except CustomError as e:
            logging.warning(str(e.value))
        except Exception as e:
            logging.warning(e)
        finally:
            self.units.dbDisconnect(conn)
            return stockSecurityData

    '''
    布林带线交易策略，根据股价与布林带上轨和布林带下轨之间的关系确定多空，在价格上穿或下穿布林带线上下轨时产生交易信号。布林带线的均线类型不可选
    1，当价格上穿上轨时，产生全仓买入信号
    2，当价格下穿下轨时，产生全仓卖出信号
    '''
    def bollinger(self, data, period=20):
        df = pd.DataFrame()
        df['bollLower'], df['bollMiddle'], df['bollUpper'] = ta.BBANDS(data.close, timeperiod=period, nbdevup=2, nbdevdn=2, matype=0)
        return df

    def ema(self, data, short=13, middle=34, long=55):
        df = pd.DataFrame()
        ta.set_compatibility(1)
        df['emaS'] = ta.EMA(data.close, short)
        df['emaM'] = ta.EMA(data.close, middle)
        df['emaL'] = ta.EMA(data.close, long)
        return df

    '''
    Volume：成交量
    '''
    def volume(self, data, short=13, long=89):
        df = pd.DataFrame()
        df['volShort'] = ta.MA(data.vol, timeperiod=short, matype=0)
        df['volLong'] = ta.MA(data.vol, timeperiod=long, matype=0)
        return df

    '''
    MACD：移动平滑异同平均线
    应用：
        1.DIF 与MACD均为正值亦即在中轴线上，大势属多头市场，DIF 向上突破 MACD，应作买。若 DIF向下跌破MACD应只可作回档，暂时获利了解。
        2.反之DIF与MACD均为负值时，即在0轴线以下时，大势属空头市场，DIF 向下跌破MACD，可作卖。若 DIF向上突破MACD只可作空头暂时补空。
        3.如同强弱指标，背离走势也适用在 MACD的图形上，当MACD图形与 K线 图趋势线发生背离时亦为反转讯号。
        4.MACD无法预知高价及低价。盘局时，失误率较高, 但如配合 RSI及KD线 应用则可以解决此二缺点.
        5.运用柱形图的变化可提早作买或作卖，免得失去一段行情，但注意有时 亦会因贪小而失大。
    从MACD的走势，投资者可以发现三种讯息： 
        1. 当MACD升穿Signal Line，入市讯息；
        2. 当MACD跌穿Signal Line，出市讯息。 
        3. 当MACD上升时，股票价格可能是超买
    交易策略:
        1.DIFF、DEA均为正，DIFF向上突破DEA，买入信号。       
        2.DIFF、DEA均为负，DIFF向下跌破DEA，卖出信号。       
        3.DEA线与K线发生背离，行情反转信号。       
        4.分析MACD柱状线，由正变负，卖出信号；由负变正，买入信号。
    '''
    def macd(self, data, short=12, long=6, signal=9):
        df = pd.DataFrame()
        df['diff'], df['dea'], df['macd'] = ta.MACD(data.close, short, long, signal)
        return df

    '''
    RSI：相对强度指数 
    交易策略：
    RSI 指标度量最近价格变化的幅度，从而判断目前股票属于超卖还是超买状态
    1, 当RSI > ulim时，设置持仓目标为1
    2, 当RSI < llim时，设置持仓目标为-1
    3, 其余情况设置持仓目标为0
    '''
    def rsi(self, data, period=14):
        df = pd.DataFrame()
        df['rsi'] = ta.RSI(data.close, period)
        return df

    '''
    CCI：商品渠道指数
    交易策略：
    CCI商品渠道指数被用来判断当前股价位于超卖还是超买区间，本策略使用这个指标生成投资仓位目标
    1, 当CCI大于0时，输出弱多头
    2, 当CCI小于0时，输出弱空头
    3, 当CCI大于50时，输出强多头
    4, 当CCI小于 - 50时，输出强空头
    '''
    def cci(self, data, period=14):
        df = pd.DataFrame()
        df['cci'] = ta.CCI(data.high, data.low, data.close, period)
        return df

    '''
    KDJ：随机指数 
    STOCH 交易策略：
    STOCH 指标度量价格变化的动量，并且动量的大小判断价格趋势，并生成比例买卖交易信号。
    1, 当k > 80时，产生逐步卖出信号，每周期卖出持有份额的30%
    2, 当k < 20时，产生逐步买入信号，每周期买入总投资额的10%
    '''
    def kdj(self, data, fastk_period=9, slowk_period=3, slowd_period=3):
        df = pd.DataFrame()
        df['K'], df['D'] = ta.STOCH(data.high, data.low, data.close, fastk_period=fastk_period, slowk_period=slowk_period, slowk_matype=1, slowd_period=slowd_period, slowd_matype=1)  # 计算kdj的正确配置
        df.loc[:, 'J'] = 3.0 * df.loc[:, 'K'] - 2.0 * df.loc[:, 'D']
        return df

    '''
    SKD：快速随机指标 
    STOCHF 交易策略：
    STOCHF 指标度量价格变化的动量，与STOCH策略类似，使用快速随机指标判断价格趋势，并生成比例买卖交易信号。
    1, 当k > 80时，产生逐步卖出信号，每周期卖出持有份额的30%
    2, 当k < 20时，产生逐步买入信号，每周期买入总投资额的10%    
    '''
    def skd(self, data, fastk_period=9, fastd_period=3):
        df = pd.DataFrame()
        df['fastk'], df['fastd'] = ta.STOCHF(data.high, data.low, data.close, fastk_period=fastk_period, fastd_period=fastd_period, fastd_matype=0)
        return df

    '''
    MFI：货币流向指数
    交易策略：
    MFI指数用于判断股价属于超买还是超卖状态，本策略使用MFI指标生成交易信号
    1, 当 MFI > 20 时，持续不断产生10 % 买入交易信号
    2, 当 MFI > 80 时，持续不断产生30 % 卖出交易信号，持续卖出持仓股票
    '''
    def mfi(self, data):
        df = pd.DataFrame()
        df['mfi'] = ta.MFI(data.high, data.low, data.close, data.vol)
        return df

    def obv(self, data):
        df = pd.DataFrame()
        obvta = ta.OBV(data['close'], data['vol'])
        obv = []
        for i in range(0, len(data)):
            if i == 0:
                obv.append(data['vol'].values[i])
            else:
                if data['close'].values[i] > data['close'].values[i - 1]:
                    obv.append(obvta[i - 1] + data['vol'].values[i])
                if data['close'].values[i] < data['close'].values[i - 1]:
                    obv.append(obvta[i - 1] - data['vol'].values[i])
                if data['close'].values[i] == data['close'].values[i - 1]:
                    obv.append(obvta[i - 1])
        df['obv'] = obv
        return df

    def atr(self, data, period=14, multiplier=1.5):
        df = pd.DataFrame()
        df['atr'] = ta.ATR(data.high, data.low, data.close, timeperiod=period)
        df['trange'] = ta.TRANGE(data.high, data.low, data.close)
        df['tr'] = ta.WMA(df['trange'], period)
        '''upper = []
        lower = []
        for i in range(0, len(data)):
            upper.append(df['tr'][i] * multiplier + data.high)
            lower.append(data.low - df['tr'][i] * multiplier)
        df['artUpper'] = upper
        df['artLower'] = lower'''
        return df

    '''
    WR：威廉指标
    WILLR 交易策略：
    WILLR 指标被用于计算股价当前处于超买还是超卖区间，并用于生成交易信号
    1, 当 WILLR > -l 时，产生逐步卖出信号，每周期卖出持有份额的30%
    2, 当 WILLR < -u 时，产生逐步买入信号，每周期买入总投资额的10%
    '''
    def willr(self, data, long=10, short=6):
        df = pd.DataFrame()
        df['willrL'] = ta.WILLR(data.high, data.low, data.close, timeperiod=long)
        df['willrS'] = ta.WILLR(data.high, data.low, data.close, timeperiod=short)
        return df

    # SAR：抛物线指标
    # acceleration：加速因子；maximum：极点价
    def sar(self, data, acceleration=2, maximum=2):
        df = pd.DataFrame()
        df['sar'] = ta.SAR(data.high, data.low, acceleration, maximum)
        df['sarext'] = ta.SAREXT(data.high, data.low)
        return df

    def cross(self, data, key1, key2, i=0):
        if(type(key1) == str):
            val1B = data.iloc[i-2][key1]
            val1A = data.iloc[i-1][key1]
        elif(type(key1) == int):
            val1B = val1A = key1
        else:
            val1B = val1A = 0
        if(type(key2) == str):
            val2B = data.iloc[i-2][key2]
            val2A = data.iloc[i-1][key2]
        elif(type(key2) == int):
            val2B = val2A = key2
        else:
            val2A = val2B = 0
        if(data.iloc[i-2][key1] < val2B and data.iloc[i-1][key1] > val2A):
            return True
        else:
            return False

class TradeStrategy(bt.Strategy):
    def log(self, txt, dt=None):
        ''' 提供记录功能'''
        dt = dt or self.datas[0].datetime.datetime(0)
        print('%s, %s' % (dt.isoformat(sep=" "), txt))

    def __init__(self):
        print('当前持仓量', self.getposition(self.data).size)
        print('当前持仓成本', self.getposition(self.data).price)
        # self.getpositionbyname(name=None, broker=None)
        for i in range(0, len(self.datas)):
            print("--------- 打印 self.datas 中 Index:" + str(i) + "，数据集名称：" + self.datas[i]._name + " ----------")

        # Add a MovingAverageSimple indicator
        '''self.emaShort = bt.talib.EMA(self.data, timeperiod=21)
        self.emaMiddle = bt.talib.EMA(self.data, timeperiod=55)
        self.emaLong = bt.talib.EMA(self.data, timeperiod=144)'''

    def next(self):
        # 目前的策略就是简单显示下收盘价。
        #self.log('Close, %.2f，Volume： %.2f, emaS：%.2f，tr：%.2f，sar：%.2f，sarext：%.2f' % (self.dataclose[0], self.datavolume[0], self.datas[0].emaS[0], self.datas[0].tr[0], self.datas[0].sar[0], self.datas[0].sarext[0]))
        #self.log('Close： %.2f，Open： %.2f，Volume： %.2f，Ema(21): %.2f，Ema(55), %.2f，Ema(144), %.2f' % (self.dataclose[0], self.dataopen[0], self.datavolume[0], self.emaShort[0], self.emaMiddle[0], self.emaLong[0]))
        for i in range(0, len(self.datas)):
            data = self.getdatabyname(self.datas[i]._name)  # 根据名称返回数据集
            if(data.macd[0] >= 0 and data.diff[-1] < data.dea[-1] and data.diff[0] > data.dea[0] and data.diff[0] > 0):
                self.log(self.datas[i]._name + ', BUY, price: ' + str(data.close[0]) + '，macd：' + str(data.macd[0]) + '，diff：' + str(data.diff[0]) + '，dea：' + str(data.dea[0]))

class StockPandasData(bt.feeds.PandasData):
    lines = ('amount', 'year', 'month', 'day', 'hour', 'minute', 'emaS', 'emaM', 'emaL', 'K', 'D', 'J', 'rsi', 'cci', 'diff', 'dea', 'macd', 'bollLower', 'bollMiddle', 'bollUpper',
             'mfi', 'obv', 'atr', 'trange', 'tr', 'willrL', 'willrS', 'sar', 'sarext')  # 要添加的线
    params = (
        ('nullvalue', 0.0),
        ('amount', 5),
        ('year', 6),
        ('month', 7),
        ('day', 8),
        ('hour', 9),
        ('minute', 10),
        ('emaS', -1),
        ('emaM', -1),
        ('emaL', -1),
        ('K', -1),
        ('D', -1),
        ('J', -1),
        ('rsi', -1),
        ('cci', -1),
        ('diff', -1),
        ('dea', -1),
        ('macd', -1),
        ('bollLower', -1),
        ('bollMiddle', -1),
        ('bollUpper', -1),
        ('mfi', -1),
        ('obv', -1),
        ('atr', -1),
        ('trange', -1),
        ('tr', -1),
        ('willrL', -1),
        ('willrS', -1),
        ('sar', -1),
        ('sarext', -1),
    )

if __name__ == '__main__':
    units = QUnits()
    setting = units.XML2ObjFromFile('./setting.xml').setting
    cerebro = bt.Cerebro()
    cerebro.addstrategy(TradeStrategy)       # 增加一个策略
    cerebro.broker.setcash(100000.0)
    # 设置固定滑点
    cerebro.broker.set_slippage_fixed(fixed=0.001)
    # 设置 0.00005 = 0.005% 的手续费
    cerebro.broker.setcommission(commission=0.00005)
    analysis = StockAnalysis()
    conn = units.dbConnect(setting.pgsql)
    try:
        cursor = conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
        cursor.execute(
            "SELECT * FROM trading.owner_stock_view  WHERE account_id = trading.func_get_default_account('GadflyBSD', 'account.stock')->'data'->>'uuid'")
        for row in cursor.fetchall():
            staff = units.getStockSecurityBars(3, row['market'], row['stock_code'], 1, 800, True)
            # print(staff.columns)
            startDateTime = datetime(staff['year'][0], staff['month'][0], staff['day'][0], staff['hour'][0],
                                     staff['minute'][0])  # 回测开始时间
            endDateTime = datetime(staff["year"].tail(1).values[0], staff["month"].tail(1).values[0],
                                   staff["day"].tail(1).values[0], staff["hour"].tail(1).values[0],
                                   staff["minute"].tail(1).values[0])
            ema = analysis.ema(staff, 21, 55, 144)
            kdj = analysis.kdj(staff)
            cci = analysis.cci(staff)
            rsi = analysis.rsi(staff)
            macd = analysis.macd(staff)
            boll = analysis.bollinger(staff)
            mfi = analysis.mfi(staff)
            obv = analysis.obv(staff)
            art = analysis.atr(staff)
            willr = analysis.willr(staff)
            sar = analysis.sar(staff)
            stocks = staff.join(ema, how='outer').join(kdj, how='outer').join(cci, how='outer').join(rsi,
                                                                                                     how='outer').join(
                macd, how='outer').join(boll, how='outer') \
                .join(mfi, how='outer').join(obv, how='outer').join(art, how='outer').join(willr, how='outer').join(sar,
                                                                                                                    how='outer')
            stocks.index = pd.to_datetime(stocks.datetime)
            stockData = StockPandasData(dataname=stocks, fromdate=startDateTime, todate=endDateTime, volume=4,
                                        openinterest=-1)
            # stockData = bt.feeds.PandasData(dataname=stocks, fromdate=startDateTime, todate=endDateTime, volume=4, openinterest=-1)
            cerebro.adddata(stockData, name=row['stock_code'])  # 将数据传入回测系统
        print('Starting Portfolio Value: %.2f' % cerebro.broker.getvalue())
        cerebro.run()
        print('Final Portfolio Value: %.2f' % cerebro.broker.getvalue())
    except CustomError as e:
        logging.warning(str(e.value))
    except Exception as e:
        logging.warning(e)
    finally:
        units.dbDisconnect(conn)
